from mem0.embeddings.accessTokenFetcher import AccessTokenFetcher
import requests
import json
from mem0.embeddings.base import EmbeddingBase

class BaiDuEmbedding(EmbeddingBase):
    def __init__(self, model="text-embedding-3-small", baidu_llm_config=None):
        self.dims = 384
        self.model = model
        self.tokenFetcher = AccessTokenFetcher(baidu_llm_config.embedder)  # 实例化AccessTokenFetcher

    def embed(self, text):
        token = self.tokenFetcher.get_access_token()  # 调用实例的get_access_token方法
        chat_url = f"https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/embeddings/embedding-v1?access_token={token}"

        ebd_request = {
            "input": [text]
        }

        headers = {'Content-Type': 'application/json'}
        response = requests.post(chat_url, data=json.dumps(ebd_request), headers=headers)

        if response.status_code == 200:
            return json.loads(response.content.decode('utf-8'))['data'][0]['embedding']
        else:
            print(f"Error: {response.status_code}, {response.text}")
            return None